package com.adnaan.backend.websocket;

import com.adnaan.backend.service.AiService;
import com.fasterxml.jackson.databind.ObjectMapper;
import lombok.Data;
import lombok.extern.slf4j.Slf4j;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;

import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;

@Slf4j
public class AiChatWebSocketHandler extends TextWebSocketHandler {

    private final ObjectMapper objectMapper = new ObjectMapper();
    private final ExecutorService executorService = Executors.newCachedThreadPool(); // 使用线程池来处理消息
    private final AiService aiService;

    public AiChatWebSocketHandler(AiService aiService) {
        this.aiService = aiService;
    }

    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) {
        executorService.submit(() -> {
            try {
                ChatRequest request = objectMapper.readValue(message.getPayload(), ChatRequest.class);
                aiService.chat(
                        request.getMessage(),
                        // 处理响应
                        content -> {
                            try {
                                Map<String, Object> wsResponse = new HashMap<>();
                                wsResponse.put("content", content);
                                wsResponse.put("role", "assistant");
                                String jsonResponse = objectMapper.writeValueAsString(wsResponse);
                                session.sendMessage(new TextMessage(jsonResponse));
                            } catch (Exception e) {
                                log.error("发送消息时出错: {}", e.getMessage());
                            }
                        },
                        // 处理错误
                        error -> handleError(session, error),
                        // 处理完成
                        () -> {
                            try {
                                session.sendMessage(new TextMessage("[DONE]"));
                            } catch (Exception e) {
                                log.error("发送已完成消息时出错: {}", e.getMessage());
                            }
                        }
                );
            } catch (Exception e) {
                log.error("处理会话消息时出错 {}: {}", session.getId(), e.getMessage());
                handleError(session, e);
            }
        });
    }

    @Override
    public void afterConnectionEstablished(WebSocketSession session) {
        log.info("AI聊天WebSocket连接已建立: {}", session.getId());
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) {
        log.info("AI聊天WebSocket连接已关闭: {}", session.getId());
    }

    @Override
    public void handleTransportError(WebSocketSession session, Throwable exception) {
        log.error("AI聊天WebSocket错误: {}", exception.getMessage());
    }

    private void handleError(WebSocketSession session, Exception error) {
        try {
            Map<String, Object> errorResponse = new HashMap<>();
            errorResponse.put("error", error.getMessage());
            session.sendMessage(new TextMessage(objectMapper.writeValueAsString(errorResponse)));
        } catch (Exception e) {
            log.error("发送错误消息时出错: {}", e.getMessage());
        }
    }

    @Data
    static class ChatRequest {
        private String message;
        private String sessionId;
    }
} 